import queue
import re
import uiautomation as auto
import threading
from pynput import keyboard
from tqdm import tqdm
import time
from datetime import datetime
import tkinter
import win32gui
import os
import sys
import csv
import openpyxl  # Add this import

from tool.node import build_node
from tool.llm import LLMClient

# 弹窗提示
def popupmsg(msg):
    # tkinter弹窗在最顶端，显示“提取已完毕”，按下“确定”后关闭窗口
    root = tkinter.Tk()
    root.title("提示")
    root.geometry("250x70")
    root.resizable(True, True)
    root.attributes("-topmost", True)
    tkinter.Label(root, text=msg).pack()
    tkinter.Button(root, text="确定", command=root.destroy).pack()
    # 定义关闭窗口的函数
    def close_window():
        root.destroy()
    # 让窗口获取焦点
    root.focus_force()
    root.bind("<Return>", lambda event: close_window())
    root.mainloop()


# 定义一个标志变量，用于控制函数是否继续运行
should_continue = True

# 定义一个队列，用于存储llm_response线程
llm_responses_queue = []

# 装饰器，用于在按下 Esc 键时终止函数运行
def esc_interruptible(stop_reminder=True):
    def decorator(func):
        def wrapper(*args, **kwargs):
            global should_continue

            def on_press(key):
                global should_continue
                try:
                    if key == keyboard.Key.esc:
                        print("你按下了 Esc 键，函数将终止。")
                        should_continue = False  # 设置标志变量为 False 以终止函数运行
                        return False  # 停止监听键盘事件
                except AttributeError:
                    pass

            # 创建监听器线程
            keyboard_listener = keyboard.Listener(on_press=on_press)
            keyboard_listener.start()
            
            # 创建目标函数线程
            def target_func():
                try:
                    func(*args, **kwargs)
                except Exception as e:
                    popupmsg(f"{e}")
                    raise e
            function_thread = threading.Thread(target=target_func)
            # function_thread = threading.Thread(target=lambda: func(*args, **kwargs))
            function_thread.start()

            try:
                # 等待目标函数线程结束
                function_thread.join()
                if stop_reminder:
                    popupmsg("程序结束")
            finally:
                # 等待键盘监听器线程结束
                keyboard_listener.stop()
                keyboard_listener.join()

        return wrapper
    return decorator


# 用于迭代子元素的生成器
def child_iterator(element):
    child = element.GetFirstChildControl()
    while child:
        yield child

        child = child.GetNextSiblingControl()

# 逐个获取订阅号的消息
@esc_interruptible()
def get_news_one_by_one(output_path, result_path, unparsable_path, MAX_ACCOUNTS=None):
    
    if MAX_ACCOUNTS is None:
        MAX_ACCOUNTS = 1000
    
    global should_continue
    
    subscription_wnd = auto.WindowControl(searchDepth=1, ClassName='SubscriptionWnd')
    # 将subscription_wnd置顶
    subscription_wnd.SetFocus()
    
    result = {}
    summary_result = []
    unparsable_result = []
    header = ['标题', '公众号', '发布日期', '阅读数', '相关程度']
    
    document_control = subscription_wnd.DocumentControl()
    account_list = document_control.GetFirstChildControl().GetFirstChildControl().GetFirstChildControl().GetFirstChildControl()
    if not account_list.Exists(0.2):
        print("Account list not found.")
        return result
    
    llm_responses_queue = []
    
    def parse_line(line):
        if line == "None":
            return None
        # 用re正则表达式严格匹配，提取标题、公众号、发布日期、阅读数、相关程度
        # 示例: "标题, 公众号, 发布日期, 阅读数, 相关程度"
        pattern = r'(?P<title>.*)\s*,\s*(?P<account>.*)\s*,\s*(?P<date>.*)\s*,\s*(?P<read_num>.*)\s*,\s*(?P<relatedness>.*)$'
        match = re.match(pattern, line)
        if match:
            return match.group('title').strip(), match.group('account').strip(), match.group('date').strip(), match.group('read_num').strip(), match.group('relatedness').strip()
        else:
            print("Failed to parse line:", line)
            unparsable_result.append(line)
            return None
    
    def check_finished_thread():
        for i, (llm_thread, response_queue) in enumerate(llm_responses_queue):
            if not llm_thread.is_alive():
                llm_responses_queue.pop(i)
                response = response_queue.get()
                for line in response.split('\n'):
                    parse_result = parse_line(line.strip())
                    if parse_result:
                        summary_result.append(parse_result)
                break
    
    def handle_exit(esc=True):
        if esc:
            print("Esc detected. Function terminated.")
        # Wait for all llm_threads to finish
        # llm_responses_queue里面的thread可以乱序结束，优先处理已经结束的thread
        if not llm_responses_queue:
            print("No LLM requests found.")
            return
        
        width = len(str(len(llm_responses_queue)))
        while llm_responses_queue:
            print("剩余LLM请求数:", f"{len(llm_responses_queue):{width}d}", end='\r')
            check_finished_thread()
        
        # # 将unparsable_result写入文件
        # if unparsable_result:
        #     with open(unparsable_path, 'w', encoding='utf-8') as f:
        #         for line in unparsable_result:
        #             f.write(line + '\n')
        #     print("The unparsable result has been saved to", unparsable_path)
        # 将summary_result写入excel
        print("Writing summary result to excel...")
        workbook = openpyxl.Workbook()
        sheet = workbook.active
        sheet.append(header)
        for row in summary_result:
            sheet.append(row)
        def set_column_width(sheet):
            # 根据内容长度的平均值自动调整列宽
            for col in sheet.columns:
                average_length = sum(len(str(cell.value)) for cell in col) / len(col)
                sheet.column_dimensions[col[0].column_letter].width = average_length * 2.2 + 10
        set_column_width(sheet)
        # 新增一个sheet写入unparsable_result
        if unparsable_result:
            sheet = workbook.create_sheet(title='无法解析的消息')
            for line in unparsable_result:
                sheet.append([line])
            set_column_width(sheet)
        workbook.save(result_path)
        workbook.close()
        print("The result has been saved to", result_path)
        try:
            # 打开excel
            os.system(f'start excel "{result_path}"')
        except Exception as e:
            print(e)
        
    
    # 检测是否按下了 Esc 键
    if not should_continue:
        handle_exit()
        return result
    
    # Scroll up and down to make sure that the content of the account list is loaded
    account_list.SetFocus()
    account_list.Click(2, 2, simulateMove=False, waitTime=0)
    # auto.SendKeys('{End}')
    auto.SendKeys('{Home}')
    first_child = account_list.GetFirstChildControl()
    first_child.Click(1, 1, simulateMove=False)
    list_length = len(account_list.GetChildren())
    for _ in range(min(MAX_ACCOUNTS, list_length - 1)):
        auto.SendKeys('{Down}', waitTime=0)
        # continue
        # 检测是否按下了 Esc 键
        if not should_continue:
            handle_exit()
            return result
    auto.SendKeys('{Home}')
    first_child = account_list.GetFirstChildControl()
    first_child.Click(1, 1, simulateMove=False)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        account_iter_num = 0
        for account_element in tqdm(child_iterator(account_list), total=list_length, disable=True):
            if account_iter_num >= MAX_ACCOUNTS:
                break
            
            # Extract information from the account element
            account = account_element.TextControl()
            account_name = account.Name
            # first_news_title = account.GetNextSiblingControl().TextControl().Name
            # print(account_name)
            
            # Dive into the middle panel to get more information
            # account_element.Click(1, 1, simulateMove=False, waitTime=0)
            middle_panel = subscription_wnd.DocumentControl().GetFirstChildControl().GetFirstChildControl().GetLastChildControl().GetFirstChildControl().GetFirstChildControl()
            
            # Wait until the middle panel is loaded
            while middle_panel.TextControl().Name != account_name:
                # 检测是否按下了 Esc 键
                if not should_continue:
                    handle_exit()
                    return result
                print(account_name, "Middle panel is not loaded.")
                account_element.Click(1, 1, simulateMove=False, waitTime=0)
                time.sleep(0.1)
            
            middle_panel_node = build_node(middle_panel, max_children=5)
            account_news = middle_panel_node.get_summary_text()
            
            # print("Start exporting middle panel to json...")
            # with open('middle_panel.json', 'w', encoding='utf-8') as f:
            #     json.dump(middle_panel_node.convert_to_dict(), f, ensure_ascii=False, indent=2)
            
            # with open('middle_panel_summary.txt', 'w', encoding='utf-8') as f:
            #     f.write(middle_panel_node.get_summary_text())
            
            # with open('middle_panel.txt', 'w', encoding='utf-8') as ff:
            #     middle_panel_node.write_node_tree(ff)
            
            result[account_name] = account_news
            
            write_str = account_name + '\n' + account_news + '\n\n'
            
            def fetch_llm_response(queue, write_str):
                response = llm_client.get_llm_response(write_str)
                queue.put(response)
            
            # 创建一个队列来存储llm_response
            response_queue = queue.Queue()
            
            # 启动一个线程来获取llm_response
            llm_thread = threading.Thread(target=fetch_llm_response, args=(response_queue, write_str))
            llm_thread.start()
            
            # 将write_str和llm_thread存储在一个队列中
            llm_responses_queue.append((llm_thread, response_queue))
            
            f.write(write_str)
            
            check_finished_thread()
            
            # 检测是否按下了 Esc 键
            if not should_continue:
                handle_exit()
                return result
            
            # Press "↓" to shift to the next account
            auto.SendKeys('{Down}', waitTime=0)
            
            account_iter_num += 1

    print("The result has been saved to", output_path)
    
    handle_exit(esc=False)
    
    return result

def show_wechat():
    '''
        Activates the window and displays it as a minimized window.
        See https://learn.microsoft.com/en-us/windows/win32/api/winuser/nf-winuser-showwindow.
    '''
    HWND = win32gui.FindWindow('WeChatMainWndForPC', None)
    win32gui.ShowWindow(HWND, 2)
    wechat_wnd = auto.WindowControl(searchDepth=1, ClassName='WeChatMainWndForPC')
    wechat_wnd.SwitchToThisWindow()

def open_subscription_wnd():
    global should_continue
    
    wechat_wnd = auto.WindowControl(searchDepth=1, ClassName='WeChatMainWndForPC')
    MainControl1 = [i for i in wechat_wnd.GetChildren() if not i.ClassName][0]
    MainControl2 = MainControl1.GetFirstChildControl()
    NavigationBox, SessionBox, ChatBox  = MainControl2.GetChildren()

    session_list = SessionBox.ListControl()
    session_list.SetFocus()
    session_list.MoveCursorToMyCenter()

    while session_list.GetFirstChildControl().BoundingRectangle.top != session_list.BoundingRectangle.top:
        if not should_continue:
            print("Esc detected. Function terminated.")
            return
        auto.WheelUp(waitTime=0)

    found = False
    while not found:
        if not should_continue:
            print("Esc detected. Function terminated.")
            return
        for item in session_list.GetChildren():
            if item.Name == "订阅号" and item.BoundingRectangle.bottom < session_list.BoundingRectangle.bottom * 2 / 3 and item.BoundingRectangle.top > session_list.BoundingRectangle.bottom / 3:
                item.DoubleClick(waitTime=0)
                found = True
                break
        if not found:
            auto.WheelDown(waitTime=0)

def wait_until_subscription_wnd_loaded():
    global should_continue
        
    subscription_wnd = auto.WindowControl(searchDepth=1, ClassName='SubscriptionWnd')
    while not subscription_wnd.Exists(3):
        # print("Waiting for the subscription window to load...")
        if not should_continue:
            print("Esc detected. Function terminated.")
            return
        pass
    sub_wnd_loaded = False
    while not sub_wnd_loaded:
        if not should_continue:
            print("Esc detected. Function terminated.")
            return
        try:
            account_list = subscription_wnd.DocumentControl().GetFirstChildControl().GetFirstChildControl().GetFirstChildControl().GetFirstChildControl()
            time.sleep(0.2)
            sub_wnd_loaded = True
        except:
            # print("Waiting for the subscription window to load...")
            pass

@esc_interruptible(stop_reminder=False)
def prepare_subscription_wnd():
    if not auto.WindowControl(searchDepth=1, ClassName='SubscriptionWnd').Exists(0.5):
        show_wechat()
        open_subscription_wnd()
    wait_until_subscription_wnd_loaded()
 
 
if __name__ == "__main__":

    auto.Logger.WriteLine = lambda *args: None  # 禁用内部日志
    auto.Logger.SetLogFile('')
    auto.SetGlobalSearchTimeout(0.5)
    
    llm_client = LLMClient()
    print("llm ready")
    
    def get_num_accounts():
        def on_submit():
            nonlocal num_accounts
            num_accounts = entry.get()
            root.destroy()

        root = tkinter.Tk()
        root.title("输入")
        root.geometry("300x130")
        root.attributes("-topmost", True)
        tkinter.Label(root, text="请输入要获取的订阅号数目:").pack(pady=10)
        entry = tkinter.Entry(root)
        entry.pack(pady=5)
        entry.focus()  # 自动focus
        tkinter.Button(root, text="确定", command=on_submit).pack(pady=10)
        root.bind("<Return>", lambda event: on_submit())  # 按回车确认

        num_accounts = None
        root.mainloop()

        return int(num_accounts) if num_accounts.isdigit() else None

    try:
        NUM_ACCOUNTS = get_num_accounts()
    except Exception as e:
        print(e)
        sys.exit(1)
    print("NUM_ACCOUNTS:", NUM_ACCOUNTS if NUM_ACCOUNTS is not None else "None")
    # NUM_ACCOUNTS = None
    # NUM_ACCOUNTS = 10
    
    minute = datetime.now().strftime("%m月%d日%H时%M分")
    target_dir = os.path.join('数据', minute)
    
    if not os.path.isdir(target_dir):
        os.makedirs(target_dir, exist_ok=True)
    # 文件名为[日期时间]_[数量]_订阅号消息.txt
    output_path = os.path.join(target_dir, "订阅号消息.txt")
    result_path = os.path.join(target_dir, "大模型总结.xlsx")
    unparsable_path = os.path.join(target_dir, "无法解析的消息.txt")
    
    subscription_wnd = auto.WindowControl(searchDepth=1, ClassName='SubscriptionWnd')
    
    if not subscription_wnd.Exists():
        print("prepare_subscription_wnd")
        prepare_subscription_wnd()
    
    if should_continue:
        get_news_one_by_one(output_path, result_path, unparsable_path, NUM_ACCOUNTS)